/**
 * Copyright (c) 2013-2024 Nikita Koksharov
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.redisson.spring.data.connection;

import org.reactivestreams.Publisher;
import org.redisson.api.StreamMessageId;
import org.redisson.client.codec.ByteArrayCodec;
import org.redisson.client.codec.StringCodec;
import org.redisson.client.protocol.RedisCommand;
import org.redisson.client.protocol.RedisCommands;
import org.redisson.client.protocol.RedisStrictCommand;
import org.redisson.reactive.CommandReactiveExecutor;
import org.springframework.data.redis.connection.ReactiveRedisConnection;
import org.springframework.data.redis.connection.ReactiveStreamCommands;
import org.springframework.data.redis.connection.stream.ByteBufferRecord;
import org.springframework.data.redis.connection.stream.RecordId;
import org.springframework.data.redis.connection.stream.StreamOffset;
import org.springframework.data.redis.connection.stream.StreamRecords;
import org.springframework.util.Assert;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.ByteBuffer;
import java.util.*;
import java.util.stream.Collectors;

/**
 *
 * @author Nikita Koksharov
 *
 */
public class RedissonReactiveStreamCommands extends RedissonBaseReactive implements ReactiveStreamCommands {

    RedissonReactiveStreamCommands(CommandReactiveExecutor executorService) {
        super(executorService);
    }

    private static List<String> toStringList(List<RecordId> recordIds) {
        return recordIds.stream().map(RecordId::getValue).collect(Collectors.toList());
    }

    @Override
    public Flux<ReactiveRedisConnection.NumericResponse<AcknowledgeCommand, Long>> xAck(Publisher<AcknowledgeCommand> publisher) {
        return execute(publisher, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getGroup(), "Group must not be null!");
            Assert.notNull(command.getRecordIds(), "recordIds must not be null!");

            List<Object> params = new ArrayList<>();
            byte[] k = toByteArray(command.getKey());
            params.add(k);
            params.add(command.getGroup());
            params.addAll(toStringList(command.getRecordIds()));

            Mono<Long> m = write(k, StringCodec.INSTANCE, RedisCommands.XACK, params.toArray());
            return m.map(v -> new ReactiveRedisConnection.NumericResponse<>(command, v));
        });
    }

    @Override
    public Flux<ReactiveRedisConnection.CommandResponse<AddStreamRecord, RecordId>> xAdd(Publisher<AddStreamRecord> publisher) {
        return execute(publisher, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getBody(), "Body must not be null!");

            byte[] k = toByteArray(command.getKey());
            List<Object> params = new LinkedList<>();
            params.add(k);

            if (!command.getRecord().getId().shouldBeAutoGenerated()) {
                params.add(command.getRecord().getId().getValue());
            } else {
                params.add("*");
            }

            for (Map.Entry<ByteBuffer, ByteBuffer> entry : command.getBody().entrySet()) {
                params.add(toByteArray(entry.getKey()));
                params.add(toByteArray(entry.getValue()));
            }

            Mono<StreamMessageId> m = write(k, StringCodec.INSTANCE, RedisCommands.XADD, params.toArray());
            return m.map(v -> new ReactiveRedisConnection.CommandResponse<>(command, RecordId.of(v.toString())));
        });
    }

    @Override
    public Flux<ReactiveRedisConnection.CommandResponse<DeleteCommand, Long>> xDel(Publisher<DeleteCommand> publisher) {
        return execute(publisher, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getRecordIds(), "recordIds must not be null!");

            byte[] k = toByteArray(command.getKey());
            List<Object> params = new ArrayList<>();
            params.add(k);
            params.addAll(toStringList(command.getRecordIds()));

            Mono<Long> m = write(k, StringCodec.INSTANCE, RedisCommands.XDEL, params.toArray());
            return m.map(v -> new ReactiveRedisConnection.CommandResponse<>(command, v));
        });
    }

    @Override
    public Flux<ReactiveRedisConnection.NumericResponse<ReactiveRedisConnection.KeyCommand, Long>> xLen(Publisher<ReactiveRedisConnection.KeyCommand> publisher) {
        return execute(publisher, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");

            byte[] k = toByteArray(command.getKey());

            Mono<Long> m = write(k, StringCodec.INSTANCE, RedisCommands.XLEN, k);
            return m.map(v -> new ReactiveRedisConnection.NumericResponse<>(command, v));
        });
    }

    @Override
    public Flux<ReactiveRedisConnection.CommandResponse<RangeCommand, Flux<ByteBufferRecord>>> xRange(Publisher<RangeCommand> publisher) {
        return range(RedisCommands.XRANGE, publisher);
    }

    private Flux<ReactiveRedisConnection.CommandResponse<RangeCommand, Flux<ByteBufferRecord>>> range(RedisCommand<?> rangeCommand, Publisher<RangeCommand> publisher) {
        return execute(publisher, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getRange(), "Range must not be null!");
            Assert.notNull(command.getLimit(), "Limit must not be null!");

            byte[] k = toByteArray(command.getKey());

            List<Object> params = new LinkedList<>();
            params.add(k);

            if (rangeCommand == RedisCommands.XRANGE) {
                params.add(command.getRange().getLowerBound().getValue().orElse("-"));
                params.add(command.getRange().getUpperBound().getValue().orElse("+"));
            } else {
                params.add(command.getRange().getUpperBound().getValue().orElse("+"));
                params.add(command.getRange().getLowerBound().getValue().orElse("-"));
            }


            if (command.getLimit().getCount() > 0) {
                params.add("COUNT");
                params.add(command.getLimit().getCount());
            }

            Mono<Map<StreamMessageId, Map<byte[], byte[]>>> m = write(k, ByteArrayCodec.INSTANCE, rangeCommand, params.toArray());
            return m.map(v -> new ReactiveRedisConnection.CommandResponse<>(command, Flux.fromStream(v.entrySet().stream()).map(e -> {
                Map<ByteBuffer, ByteBuffer> map = e.getValue().entrySet().stream()
                        .collect(Collectors.toMap(entry -> ByteBuffer.wrap(entry.getKey()),
                                entry -> ByteBuffer.wrap(entry.getValue())));
                return StreamRecords.newRecord()
                        .in(command.getKey())
                        .withId(RecordId.of(e.getKey().toString()))
                        .ofBuffer(map);
            })));
        });
    }

    @Override
    public Flux<ReactiveRedisConnection.CommandResponse<ReadCommand, Flux<ByteBufferRecord>>> read(Publisher<ReadCommand> publisher) {
        return execute(publisher, command -> {

            Assert.notNull(command.getStreamOffsets(), "StreamOffsets must not be null!");
            Assert.notNull(command.getReadOptions(), "ReadOptions must not be null!");

            List<Object> params = new ArrayList<>();

            if (command.getConsumer() != null) {
                params.add("GROUP");
                params.add(command.getConsumer().getGroup());
                params.add(command.getConsumer().getName());
            }

            if (command.getReadOptions().getCount() != null && command.getReadOptions().getCount() > 0) {
                params.add("COUNT");
                params.add(command.getReadOptions().getCount());
            }

            if (command.getReadOptions().getBlock() != null && command.getReadOptions().getBlock() > 0) {
                params.add("BLOCK");
                params.add(command.getReadOptions().getBlock());
            }

            if (command.getConsumer() != null && command.getReadOptions().isNoack()) {
                params.add("NOACK");
            }

            params.add("STREAMS");
            for (StreamOffset<ByteBuffer> streamOffset : command.getStreamOffsets()) {
                params.add(toByteArray(streamOffset.getKey()));
            }

            for (StreamOffset<ByteBuffer> streamOffset : command.getStreamOffsets()) {
                params.add(streamOffset.getOffset().getOffset());
            }

            Mono<Map<String, Map<StreamMessageId, Map<byte[], byte[]>>>> m;

            if (command.getConsumer() == null) {
                if (command.getReadOptions().getBlock() != null && command.getReadOptions().getBlock() > 0) {
                    m = read(toByteArray(command.getStreamOffsets().get(0).getKey()), ByteArrayCodec.INSTANCE, executorService.getServiceManager().getXReadBlockingCommand(), params.toArray());
                } else {
                    m = read(toByteArray(command.getStreamOffsets().get(0).getKey()), ByteArrayCodec.INSTANCE, executorService.getServiceManager().getXReadCommand(), params.toArray());
                }
            } else {
                if (command.getReadOptions().getBlock() != null && command.getReadOptions().getBlock() > 0) {
                    m = read(toByteArray(command.getStreamOffsets().get(0).getKey()), ByteArrayCodec.INSTANCE, executorService.getServiceManager().getXReadGroupBlockingCommand(), params.toArray());
                } else {
                    m = read(toByteArray(command.getStreamOffsets().get(0).getKey()), ByteArrayCodec.INSTANCE, executorService.getServiceManager().getXReadGroupCommand(), params.toArray());
                }
            }

            return m.map(v -> new ReactiveRedisConnection.CommandResponse<>(command, Flux.fromStream(v.entrySet().stream())
                    .map(ee -> {
                        return ee.getValue().entrySet().stream().map(e -> {
                            Map<ByteBuffer, ByteBuffer> map = e.getValue().entrySet().stream()
                                    .collect(Collectors.toMap(entry -> ByteBuffer.wrap(entry.getKey()),
                                            entry -> ByteBuffer.wrap(entry.getValue())));
                            return StreamRecords.newRecord()
                                    .in(ee.getKey())
                                    .withId(RecordId.of(e.getKey().toString()))
                                    .ofBuffer(map);
                        });
                    }).flatMap(Flux::fromStream)
            ));

        });
    }

    private static final RedisStrictCommand<String> XGROUP_STRING = new RedisStrictCommand<>("XGROUP");

    @Override
    public Flux<ReactiveRedisConnection.CommandResponse<GroupCommand, String>> xGroup(Publisher<GroupCommand> publisher) {
        return execute(publisher, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getGroupName(), "GroupName must not be null!");

            byte[] k = toByteArray(command.getKey());

            if (command.getAction().equals(GroupCommand.GroupCommandAction.CREATE)) {
                Assert.notNull(command.getReadOffset(), "ReadOffset must not be null!");

                Mono<String> m = write(k, StringCodec.INSTANCE, XGROUP_STRING, "CREATE", k, command.getGroupName(), command.getReadOffset().getOffset(), "MKSTREAM");
                return m.map(v -> new ReactiveRedisConnection.CommandResponse<>(command, v));
            }

            if (command.getAction().equals(GroupCommand.GroupCommandAction.DELETE_CONSUMER)) {
                Assert.notNull(command.getConsumerName(), "ConsumerName must not be null!");

                Mono<Long> m = write(k, StringCodec.INSTANCE, RedisCommands.XGROUP_LONG, "DELCONSUMER", k, command.getGroupName(), command.getConsumerName());
                return m.map(v -> new ReactiveRedisConnection.CommandResponse<>(command, v > 0 ? "OK" : "Error"));
            }

            if (command.getAction().equals(GroupCommand.GroupCommandAction.DESTROY)) {
                Mono<Long> m = write(k, StringCodec.INSTANCE, RedisCommands.XGROUP_LONG, "DESTROY", k, command.getGroupName());
                return m.map(v -> new ReactiveRedisConnection.CommandResponse<>(command, v > 0 ? "OK" : "Error"));
            }

            throw new IllegalArgumentException("unknown command " + command.getAction());
        });
    }

    @Override
    public Flux<ReactiveRedisConnection.CommandResponse<RangeCommand, Flux<ByteBufferRecord>>> xRevRange(Publisher<RangeCommand> publisher) {
        return range(RedisCommands.XREVRANGE, publisher);
    }

    @Override
    public Flux<ReactiveRedisConnection.NumericResponse<ReactiveRedisConnection.KeyCommand, Long>> xTrim(Publisher<TrimCommand> publisher) {
        return execute(publisher, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getCount(), "Count must not be null!");

            byte[] k = toByteArray(command.getKey());

            Mono<Long> m = write(k, StringCodec.INSTANCE, RedisCommands.XTRIM, k, "MAXLEN", command.getCount());
            return m.map(v -> new ReactiveRedisConnection.NumericResponse<>(command, v));
        });
    }
}
